feat(orchestration): give services on a shared node their own GPUs by wprazuch · Pull Request #3870 · NVIDIA-NeMo/Gym

wprazuch · 2026-10-01T11:44:37Z

Follow-up to the #3654 review (#3654 (comment)): automatic GPU placement, agreed to land separately. Replaces #3707 (closed), rebased onto main.

The problem

Every service on a node sees every GPU. A 1-GPU judge and a 2-GPU policy on one 4-GPU node both start on GPU 0. Today the only fix is a hand-written CUDA_VISIBLE_DEVICES per service.

What changed

Own GPUs per service. When several GPU services share a node, each vLLM service gets its own CUDA_VISIBLE_DEVICES (TP x PP x instances), lowest free GPUs first, in declaration order. A Ray service gets the leftovers.
Exported inside the step. Slurm resets an env K=V srun value to all GPUs; an export inside the step survives. A manual CUDA_VISIBLE_DEVICES now takes effect too, and its GPUs count as taken.
Opt out with share_gpus: true. Overlapping manual sets or over-subscription fail at config validation with a clear error.
nodes: auto. A pool sized by Gym: the policy goes on the first node, other services biggest-first into the first node with room, never split.

compute:
  cluster:
    node_pools:
      gpu: {partition: batch, nodes: auto, gpus_per_node: 4}
services:
  judge:  {type: vllm, model: /judge}                           # -> GPU 0
  policy: {type: vllm, model: /policy, tensor_parallel_size: 2} # -> GPUs 1,2
# -> one node

What stays the same

A node with a single GPU service renders nothing new. Explicit node counts are unchanged.
vllm_pd: a pool holding a prefill/decode tier is never packed, the router takes no GPUs, and the job script is byte-identical to main. nodes: auto refuses a P/D pool.

Out of scope

Ray-placed multi-node services, an unpinned vLLM service in a multi-node job, nodes: auto for P/D pools, and moving services between explicitly sized pools.

Testing

Unit tests in tests/unit_tests/test_slurm_script.py: disjoint sets, manual wins, errors, opt-out, Ray leftovers, per-pool packing, in-step export, nodes: auto, and vllm_pd byte-identity. All 495 orchestration tests pass; ruff clean.
Mutation checks: breaking disjointness, packing beside a P/D tier, counting the router as a GPU user, letting nodes: auto take a P/D pool, or probing everything on localhost each fails at least one test.

Cluster evidence (at `7848acfe4`, before a conflict-free rebase onto main; oci-hsg, 4-GPU GB200, `gym eval submit`, Qwen3-4B, AIME25 subset)

Job	Config	Checked	Result
7612233	`nodes: 1`, policy TP2 + judge TP1	Policy on GPUs 0,1, judge on GPU 2, GPU 3 empty (`nvidia-smi`). 96/96 rollouts, no OOM.	pass
7612235	`nodes: auto`, 6 GPUs of need	Resolved to 2 nodes; the off-node model passed its probe and served. 96/96 rollouts.	pass
7612237	`vllm_pd`, 1 prefill + 1 decode node	Byte-identical to main; router and both tiers healthy. 32/32 rollouts, no KV transfer errors.	pass

copy-pr-bot · 2026-10-01T11:44:41Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

When two or more vLLM services share a node, each now gets its own CUDA_VISIBLE_DEVICES (TP x PP x instances per node, in declaration order). A ray service there gets the leftover GPUs. A node with one GPU service renders nothing new. A manual CUDA_VISIBLE_DEVICES index list wins and reserves its GPUs. Overlaps and over-subscription fail validation. `share_gpus: true` opts a service out. Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

…vLLM beside Ray Slurm resets CUDA_VISIBLE_DEVICES to the step's GPUs, so a value passed as `env K=V srun` never reaches the task. Checked on a 4-GPU GB200 node: the env-prefix value came back as 0,1,2,3, while `export` inside the step's command left CUDA seeing only the planned GPUs. Planned and manual values on a packed node are now exported in the step. A lone vLLM service beside a GPU-using ray service is now packed too, so Ray no longer advertises the model's GPUs. A ray service with num_gpus: 0, or a single GPU service on a node, still renders nothing new. Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

A pool with `nodes: auto` is sized by Gym. The policy goes on the pool's first node, beside the driver. The other vLLM services go biggest first into the first node with room, never split. A service bigger than one node gets ceil(need / gpus_per_node) whole nodes of its own. More than one node splits the pool into sub-pools `<pool>-<i>`, placed with the existing node-pool --nodelist mechanism; a ray service spanning the pool spans them all and gets each node's leftover GPUs. Explicit node counts are unchanged. Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

A service's own CUDA_VISIBLE_DEVICES was passed as `env K=V srun`, which Slurm resets to the step's GPUs, so it was silently ignored. It is now exported inside the step's command on every node, including a node with a single service. A runtime:VAR value is relayed as GYM_CUDA_VISIBLE_DEVICES so the step still exports the batch shell's $VAR. Drop runtime:CUDA_VISIBLE_DEVICES from the multi-instance example: it re-exported the job's full GPU set and did nothing. Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

A prefill/decode tier owns its nodes, so a pool holding one is never packed, and the vllm-router takes no GPUs. `nodes: auto` refuses a pool a vllm_pd service uses: P/D pools need explicit node counts for now. Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

… lands on Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

prokotg

We discussed that information about gpus needed should flow from services and pools should be user-facing concepts with no pools modified.

There might be a problem with how pd vllm service is placed but we consider it outside of the scope of this PR so anything is acceptable

prokotg · 2026-10-05T12:37:23Z

 class NodePool(_StrictModel):
    partition: str
-    nodes: int = 1
+    # "auto": Gym packs the pool's vLLM services onto as few nodes as fit (see _resolve_auto_pool).


That's pretty slick but I want to highlight that this option might blow out when someone makes a mistake in services. It's not to say we shouldn't have auto mode but that we accept the risk

prokotg · 2026-10-05T12:41:36Z

+            "size it by. Pin a service to it with node_pool, or give it an explicit node count."
+        )
+
+    def need(service: VllmServiceConfig) -> int:


need what? :)

prokotg · 2026-10-05T12:54:58Z

+def _manual_gpus(service: BaseServiceConfig) -> list[int] | None:
+    value = service.env.get(_CUDA_VISIBLE_DEVICES)
+    if value is None or not _GPU_INDEX_LIST_RE.match(value):
+        return None
+    return [int(index) for index in value.split(",")]


Question about assumptions: how would that work for asymmetrical multi-node jobs ? do we assume jobs will be symmetrical, meaning each node assigned to the job will use the same number of gpus?

prokotg · 2026-10-05T13:04:03Z

+def _gpus_needed_per_node(service: VllmServiceConfig, nodes: int, gpus_per_node: int) -> int | None:
+    """GPUs this service uses on each of its nodes, or None when Ray places it across nodes."""
+    tp_pp = service.tensor_parallel_size * service.pipeline_parallel_size
+    if nodes <= 1:
+        return tp_pp * service.number_of_instances
+    if service.number_of_instances > 1 and not effective_ray_serve(service, nodes, [gpus_per_node]):
+        return tp_pp * (service.number_of_instances // nodes)
+    return None


In my view, this function should be per-service-class property and the process that plans gpus should read each service property to make a call

prokotg · 2026-10-05T13:07:02Z

+def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:
+    """The CUDA_VISIBLE_DEVICES to render, keyed by node pool and then service.
+
+    A pool is planned only when a vLLM service on it shares it with another service that wants


I know vLLM service began discussion about the automatic placement but this might be confusing for anyone reading. IIUC this should be service-agnostic

prokotg · 2026-10-05T15:00:05Z

                    # vLLM doesn't require auth; dummy key satisfies clients that require the header.
                    benchmark.run["policy_api_key"] = "dummy"  # pragma: allowlist secret

+        plan_gpus(self)


prokotg · 2026-10-07T08:28:35Z

+    return None
+
+
+def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:


I think plan_gpus should live outside of API

prokotg · 2026-10-07T08:30:31Z

+    The policy goes first, on the pool's first node, where the driver reaches it on localhost.
+    The rest go biggest first into the first node with room. A service bigger than one node
+    gets ceil(need / gpus_per_node) whole nodes to itself. More than one node splits the pool
+    into sub-pools named `<pool>-<i>`, which the rest of the executor places like any pool.


Compute Pool should be a user-facing concept and we should not alter it. It's a declaration of compute user has. Executors should learn how to place services within a pool.

prokotg · 2026-10-07T09:10:01Z

+    for pool_name, pool in pools.items():
+        if pool.gpus_per_node is None:
+            continue
+        on_pool = {
+            name: service
+            for name, service in config.deployed_services.items()
+            if single_node
+            or (
+                pool_name in (*service.node_pools, driver_pool)
+                if isinstance(service, RayServiceConfig)
+                else service.node_pool == pool_name
+            )
+        }
+        if any(isinstance(s, VllmPDTierConfig) for s in on_pool.values()):
+            continue
+        claimants = {
+            n: s
+            for n, s in on_pool.items()
+            if isinstance(s, (VllmServiceConfig, RayServiceConfig)) and _takes_own_gpus(s)
+        }
+        vllms = {n: s for n, s in claimants.items() if isinstance(s, VllmServiceConfig)}
+        if not vllms or len(claimants) < 2:
+            continue
+        needs = {n: _gpus_needed_per_node(s, pool.nodes, pool.gpus_per_node) for n, s in vllms.items()}
+        if None in needs.values():
+            continue
+        manual = {n: gpus for n, s in claimants.items() if (gpus := _manual_gpus(s)) is not None}
+        needs.update({n: len(gpus) for n, gpus in manual.items()})


Can I offer to simplify gpu request discovery? Instead of differentiating on different services within an external function based on subpools, I believe the information about the service resource shape and number should come directly from the service:

A property gpus_needed on BaseServiceConfig replaces the vLLM-specific logic:

0 for share_gpus or non-GPU services

len(manual) for a hand-set CUDA_VISIBLE_DEVICES

otherwise the subclass's count (for vLLM, TP × PP × instances)

place_services replaces plan_gpus and does both node-level and GPU-level packing, without isinstance(VllmServiceConfig) checks.

prokotg · 2026-10-07T09:11:29Z

    partition: str
-    nodes: int = 1
+    # "auto": Gym packs the pool's vLLM services onto as few nodes as fit (see _resolve_auto_pool).
+    nodes: int | Literal["auto"] = 1


This comes directly from my comment on sub-pools but "auto" should resolve to cumulative number of nodes

wprazuch marked this pull request as ready for review October 2, 2026 09:44

nemo-automation-bot Bot requested review from a team October 2, 2026 09:45

copy-pr-bot Bot temporarily deployed to public October 2, 2026 09:45 Inactive

copy-pr-bot Bot temporarily deployed to public October 2, 2026 09:47 Inactive

wprazuch added 6 commits October 2, 2026 11:47

test(orchestration): probe a service nodes: auto moves on the node it…

62ecf03

… lands on Signed-off-by: Wojciech Prazuch <wprazuch@nvidia.com>

copy-pr-bot Bot temporarily deployed to public October 2, 2026 09:48 Inactive

wprazuch force-pushed the wprazuch/auto-gpu-packing-v2 branch from 7848acf to 62ecf03 Compare October 2, 2026 09:48

copy-pr-bot Bot temporarily deployed to public October 2, 2026 09:50 Inactive

copy-pr-bot Bot temporarily deployed to public October 2, 2026 09:52 Inactive

copy-pr-bot Bot temporarily deployed to public October 2, 2026 09:53 Inactive

github-actions Bot added the sla:triage-overdue Review assignment is over the one-business-day SLA label Oct 5, 2026

github-actions Bot mentioned this pull request Oct 5, 2026

[Tracker] Pull request handoff SLA #3402

Open

wprazuch mentioned this pull request Oct 6, 2026

feat(orchestration): auto-size prefill/decode tier pools with nodes: auto #4094

Open

prokotg requested changes Oct 7, 2026

View reviewed changes

github-actions Bot added sla:author-overdue Author response is over the one-business-day SLA and removed sla:triage-overdue Review assignment is over the one-business-day SLA labels Oct 7, 2026

github-actions Bot mentioned this pull request Oct 22, 2026

[Tracker] Pull request handoff SLA #4925

Open

github-actions Bot mentioned this pull request Oct 22, 2026

[Tracker] Pull request handoff SLA #4926

Open

Navigation Menu

Repository navigation

feat(orchestration): give services on a shared node their own GPUs - #3870

feat(orchestration): give services on a shared node their own GPUs#3870
wprazuch wants to merge 6 commits into
mainfrom
wprazuch/auto-gpu-packing-v2

wprazuch commented Oct 1, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Oct 1, 2026

Uh oh!

prokotg left a comment

Uh oh!

prokotg Oct 5, 2026

Uh oh!

prokotg Oct 5, 2026

Uh oh!

prokotg Oct 5, 2026

Uh oh!

prokotg Oct 5, 2026

Uh oh!

prokotg Oct 5, 2026

Uh oh!

prokotg Oct 5, 2026

Uh oh!

prokotg Oct 7, 2026

Uh oh!

prokotg Oct 7, 2026

Uh oh!

prokotg Oct 7, 2026

Uh oh!

prokotg Oct 7, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return None


		def plan_gpus(config: "SubmitConfig") -> dict[str, dict[str, str]]:

Conversation

wprazuch commented Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

The problem

What changed

What stays the same

Out of scope

Testing

Cluster evidence (at 7848acfe4, before a conflict-free rebase onto main; oci-hsg, 4-GPU GB200, gym eval submit, Qwen3-4B, AIME25 subset)

Uh oh!

copy-pr-bot Bot commented Oct 1, 2026

Uh oh!

prokotg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wprazuch commented Oct 1, 2026 •

edited

Loading

Cluster evidence (at `7848acfe4`, before a conflict-free rebase onto main; oci-hsg, 4-GPU GB200, `gym eval submit`, Qwen3-4B, AIME25 subset)